# 主要目标: 按月度分析哪个月应该买入，哪个月应该卖出股票数据
# 日志
# 2025/01/28 新建本脚本，并初步实现基本功能。

import baostock as bs
import pandas as pd
from datetime import datetime
import os,sys
import matplotlib.pyplot as plt
import seaborn as sns

def download(csvfile):
    #### 登陆系统 ####
    lg = bs.login()
    # 显示登陆返回信息
    print('login respond error_code:'+lg.error_code)
    print('login respond  error_msg:'+lg.error_msg)

    #### 获取沪深A股历史K线数据 ####
    # 详细指标参数，参见“历史行情指标参数”章节；“分钟线”参数与“日线”参数不同。“分钟线”不包含指数。
    # 分钟线指标：date,time,code,open,high,low,close,volume,amount,adjustflag
    # 周月线指标：date,code,open,high,low,close,volume,amount,adjustflag,turn,pctChg
    rs = bs.query_history_k_data_plus("sh.000001",
        "date,code,open,high,low,close",
        start_date='2004-12-01', end_date='2025-12-31',
        frequency="m", adjustflag="3")
    print('query_history_k_data_plus respond error_code:'+rs.error_code)
    print('query_history_k_data_plus respond  error_msg:'+rs.error_msg)
    data_list = []
    while (rs.error_code == '0') & rs.next():
        # 获取一条记录，将记录合并在一起
        data_list.append(rs.get_row_data())
    bs.logout()
    print('logout')

    # 定义列的数据类型
    dtype_dict = {
        'code': 'object',  # 字符串类型
        'open': 'float64', # 浮点型
        'high': 'float64', # 浮点型
        'low': 'float64',  # 浮点型
        'close': 'float64' # 浮点型
    }
    df = pd.DataFrame(data_list, columns=rs.fields)#,dtype=dtype_dict,parse_dates=['date'])
    # 转换列的数据类型
    df['date'] = pd.to_datetime(df['date'])  # 将 date 列转换为日期型
    df['code'] = df['code'].astype('object')  # 将 code 列转换为字符串类型
    df[['open', 'high', 'low', 'close']] = df[['open', 'high', 'low', 'close']].astype('float64')  # 将其他列转换为浮点型
    df['preclose']=df['close']
    df['preclose'] = df['preclose'].shift(1)
    df = df.drop(0)
    df['increment'] = df['close'] - df['preclose']
    df['year'] = df['date'].dt.year  # 提取年份
    df['month'] = df['date'].dt.month  # 提取月份
    df.to_csv(csvfile, index=False)

#### 结果集输出到csv文件 ####   
csvfile = "sh.000001_month_kdata.csv"
if True or not os.path.exists(csvfile):
    download(csvfile)

# 1. 加载数据
df = pd.read_csv(csvfile)

print(df.head())


monthly_summary = df.groupby(['year', 'month'])['increment'].sum().reset_index()
print(monthly_summary)



# 创建一个以年为纵坐标，月为横坐标的表格
pivot_table = df.pivot(index='year', columns='month', values='increment')

# 打印表头（月份）
print("Year |", " | ".join(f"{month:^7}" for month in pivot_table.columns))
print("-" * 124)
# 打印每一年的数据
for year, row in pivot_table.iterrows():
    # 格式化每一行的数据
    row_values = [f"{value:7.1f}" if not pd.isna(value) else "  NA" for value in row]
    print(f"{year:4} |", " | ".join(row_values))
print("-" * 124)


#sys.exit()
# 5. 创建透视表，用于绘图
pivot_table = monthly_summary.pivot(index='year', columns='month', values='increment')


# 6. 绘制热力图
# 设置全局字体为支持中文的字体
plt.rcParams['font.sans-serif'] = ['SimHei']  # Windows 系统
# plt.rcParams['font.sans-serif'] = ['Arial Unicode MS']  # macOS 系统
# plt.rcParams['font.sans-serif'] = ['WenQuanYi Zen Hei']  # Linux 系统
# 解决负号显示为方块的问题
plt.rcParams['axes.unicode_minus'] = False
plt.figure(figsize=(12, 6))
sns.heatmap(pivot_table, annot=True, fmt=".1f", cmap='coolwarm', linewidths=0.5)
plt.title('上证月度收益 (月收盘-上月收盘) 年月分析图')
plt.xlabel('月份')
plt.ylabel('年度')
plt.show()
 